/*
 * Copyright (c) 2011 Seoul National University
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met: redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer;
 * redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution;
 * neither the name of the copyright holders nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * Authors: Jung Ho Ahn  ( gajh@snu.ac.kr )
 *
 */

#include <cstdlib>
#include <cstring>
#include "Kmer.h"

using namespace std;

int32_t Kmer::lenQ = 0;
int32_t Kmer::lenR = 0;


ostream & operator<<(ostream &out, nucleotide_type n)
{
  switch (n)
  {
    case nt_A : out << "A"; break;
    case nt_T : out << "T"; break;
    case nt_G : out << "G"; break;
    case nt_C : out << "C"; break;
    default   : out << "."; break;
  }
  return out;
}


ostream & operator<<(ostream &out, const Kmer &k)
{
  Kmer copy(k);
  string s;

  for (int32_t i = 0; i < Kmer::lenR + nNtPerUnit*Kmer::lenQ; i++)
  {
    nucleotide_type pop = copy.popLSB();
    switch (pop)
    {
      case nt_A : s = "A" + s; break;
      case nt_T : s = "T" + s; break;
      case nt_G : s = "G" + s; break;
      case nt_C : s = "C" + s; break;
      default   : s = "." + s; break;
    }
  }
  //out << "[kmer:" << s << "]";
  out << s;
  return out;
}


#if defined __MAX_KMER_SIZE_31__

Kmer & Kmer::operator=(const Kmer & k)
{
  kmer = k.kmer;
  return *this;
}


nucleotide_type Kmer::popLSB()
{
  //KmerUnit carry = 0;

  KmerUnit carry_backup = kmer & 0x3;
  kmer  = (kmer >> 2);
  //kmer |= carry << nBitKmerUnitMinusNt;
  //carry = carry_backup;
  return (nucleotide_type) carry_backup;
}


void Kmer::pushFromLSB(nucleotide_type n)
{
  KmerUnit carry        = n;

  kmer  = (kmer << 2) | carry;
  kmer &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);
}


void Kmer::pushFromMSB(nucleotide_type n)
{
  // gajh assumes that length is an odd number
  assert(lenR != 0);
  popLSB();
  kmer |= (((KmerUnit)n) << ((lenR - 1) << 1));
}


void Kmer::reverseComplement()
{
  kmer  = ~kmer;
  kmer &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);

  int32_t currHR = lenR - 1;
  int32_t currLR = 0;

  while (currHR != currLR)
  {
    uint64_t currHN = (kmer >> (currHR << 1)) & 0x3;
    uint64_t currLN = (kmer >> (currLR << 1)) & 0x3;
    kmer = (kmer & (~((KmerUnit)0x3 << (currHR << 1)))) + (currLN << (currHR << 1));
    kmer = (kmer & (~((KmerUnit)0x3 << (currLR << 1)))) + (currHN << (currLR << 1));

    currHR--;
    currLR++;
  }
}


#elif defined __MAX_KMER_SIZE_63__

Kmer & Kmer::operator=(const Kmer & k)
{
  kmerL = k.kmerL;
  kmerH = k.kmerH;
  return *this;
}


nucleotide_type Kmer::popLSB()
{
  KmerUnit carry = kmerH & 0x3;
  KmerUnit carry_backup = kmerL & 0x3;

  kmerH  = (kmerH >> 2);
  kmerL  = (kmerL >> 2);
  kmerL |= carry << nBitKmerUnitMinusNt;

  return (nucleotide_type) carry_backup;
}


void Kmer::pushFromLSB(nucleotide_type n)
{
  KmerUnit carry        = n;
  KmerUnit carry_backup = (kmerL & filter) >> nBitKmerUnitMinusNt;

  kmerL  = (kmerL << 2) | carry;
  kmerH  = (kmerH << 2) | carry_backup;
  kmerH &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);
}


void Kmer::pushFromMSB(nucleotide_type n)
{
  // gajh assumes that length is an odd number
  assert(lenR != 0);
  popLSB();
  kmerH |= (((KmerUnit)n) << ((lenR - 1) << 1));
}


void Kmer::reverseComplement()
{
  kmerL = ~kmerL;
  kmerH = ~kmerH;
  kmerH &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);

  int32_t currHR = lenR - 1;
  int32_t currLR = 0;

  while (true)
  {
    uint64_t currHN = (kmerH >> (currHR << 1)) & 0x3;
    uint64_t currLN = (kmerL >> (currLR << 1)) & 0x3;
    kmerH = (kmerH & (~((KmerUnit)0x3 << (currHR << 1)))) + (currLN << (currHR << 1));
    kmerL = (kmerL & (~((KmerUnit)0x3 << (currLR << 1)))) + (currHN << (currLR << 1));

    if (currHR == 0)
    {
      currHR = nNtPerUnit - 1;
      currLR++;
      break;
    }
    else
    {
      currHR--;
    }
    currLR++;
  }
  while (currHR != currLR)
  {
    uint64_t currHN = (kmerL >> (currHR << 1)) & 0x3;
    uint64_t currLN = (kmerL >> (currLR << 1)) & 0x3;
    kmerL = (kmerL & (~((KmerUnit)0x3 << (currHR << 1)))) + (currLN << (currHR << 1));
    kmerL = (kmerL & (~((KmerUnit)0x3 << (currLR << 1)))) + (currHN << (currLR << 1));

    currHR--;
    currLR++;
  }
}


#elif defined __MAX_KMER_SIZE_95__

Kmer & Kmer::operator=(const Kmer & k)
{
  kmerL = k.kmerL;
  kmerM = k.kmerM;
  kmerH = k.kmerH;
  return *this;
}


nucleotide_type Kmer::popLSB()
{
  KmerUnit carry = kmerH & 0x3;

  kmerH  = (kmerH >> 2);
  KmerUnit carry_backup = kmerM & 0x3;
  kmerM  = (kmerM >> 2);
  kmerM |= carry << nBitKmerUnitMinusNt;
  carry  = carry_backup;
  carry_backup = kmerL & 0x3;
  kmerL  = (kmerL >> 2);
  kmerL |= carry << nBitKmerUnitMinusNt;
  return (nucleotide_type) carry_backup;
}


void Kmer::pushFromLSB(nucleotide_type n)
{
  KmerUnit carry        = n;
  KmerUnit carry_backup = (kmerL & filter) >> nBitKmerUnitMinusNt;

  kmerL  = (kmerL << 2) | carry;
  carry  = carry_backup;
  carry_backup = (kmerM & filter) >> nBitKmerUnitMinusNt;
  kmerM  = (kmerM << 2) | carry;
  kmerH  = (kmerH << 2) | carry_backup;
  kmerH &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);
}


void Kmer::pushFromMSB(nucleotide_type n)
{
  // gajh assumes that length is an odd number
  assert(lenR != 0);
  popLSB();
  kmerH |= (((KmerUnit)n) << ((lenR - 1) << 1));
}


void Kmer::reverseComplement()
{
  kmerL = ~kmerL;
  kmerM = ~kmerM;
  kmerH = ~kmerH;
  kmerH &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);

  int32_t currHR = lenR - 1;
  int32_t currLR = 0;

  while (true)
  {
    uint64_t currHN = (kmerH >> (currHR << 1)) & 0x3;
    uint64_t currLN = (kmerL >> (currLR << 1)) & 0x3;
    kmerH = (kmerH & (~((KmerUnit)0x3 << (currHR << 1)))) + (currLN << (currHR << 1));
    kmerL = (kmerL & (~((KmerUnit)0x3 << (currLR << 1)))) + (currHN << (currLR << 1));

    if (currHR == 0)
    {
      currHR = nNtPerUnit - 1;
      currLR++;
      break;
    }
    else
    {
      currHR--;
    }
    currLR++;
  }
  while (true)
  {
    uint64_t currHN = (kmerM >> (currHR << 1)) & 0x3;
    uint64_t currLN = (kmerL >> (currLR << 1)) & 0x3;
    kmerM = (kmerM & (~((KmerUnit)0x3 << (currHR << 1)))) + (currLN << (currHR << 1));
    kmerL = (kmerL & (~((KmerUnit)0x3 << (currLR << 1)))) + (currHN << (currLR << 1));

    currHR--;

    if (currLR == nNtPerUnit - 1)
    {
      currLR = 0;
      break;
    }
    else
    {
      currLR++;
    }
  }
  while (currHR != currLR)
  {
    uint64_t currHN = (kmerM >> (currHR << 1)) & 0x3;
    uint64_t currLN = (kmerM >> (currLR << 1)) & 0x3;
    kmerM = (kmerM & (~((KmerUnit)0x3 << (currHR << 1)))) + (currLN << (currHR << 1));
    kmerM = (kmerM & (~((KmerUnit)0x3 << (currLR << 1)))) + (currHN << (currLR << 1));

    currHR--;
    currLR++;
  }
}


#else

Kmer::Kmer()
{
  kmer = new KmerUnit[lenQ + 1];

  for (int32_t i = 0; i <= lenQ; i++)
  {
    kmer[i] = 0;
  }
}


Kmer::Kmer(const Kmer& k)
{
  kmer = new KmerUnit[lenQ + 1];

  for (int32_t i = 0; i <= lenQ; i++)
  {
    kmer[i] = k.kmer[i];
  }
}


Kmer & Kmer::operator=(const Kmer & k)
{
  for (int32_t i = 0; i <= lenQ; i++)
  {
    kmer[i] = k.kmer[i];
  }
  return *this;
}


Kmer::~Kmer()
{
  /*if (kmer != NULL)
  {
    delete [] kmer;
  }*/
}


nucleotide_type Kmer::popLSB()
{
  KmerUnit carry = 0;

  for (int32_t i = lenQ; i >= 0; i--)
  {
    KmerUnit carry_backup = kmer[i] & 0x3;
    kmer[i]  = (kmer[i] >> 2);
    kmer[i] |= carry << nBitKmerUnitMinusNt;
    carry    = carry_backup;
  }
  return (nucleotide_type) carry;
}


void Kmer::pushFromLSB(nucleotide_type n)
{
  KmerUnit carry        = n;

  for (int32_t i = 0; i < lenQ; i++)
  {
    KmerUnit carry_backup = (kmer[i] & filter) >> nBitKmerUnitMinusNt;
    kmer[i] = (kmer[i] << 2) | carry;
    carry   = carry_backup;
  }

  kmer[lenQ]  = (kmer[lenQ] << 2) | carry;
  kmer[lenQ] &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);
}


void Kmer::pushFromMSB(nucleotide_type n)
{
  // gajh assumes that length is an odd number
  assert(lenR != 0);
  popLSB();
  kmer[lenQ] |= (((KmerUnit)n) << ((lenR - 1) << 1));
  //kmer[lenQ] &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);
}


void Kmer::reverseComplement()
{
  for (int32_t i = 0; i <= lenQ; i++)
  {
    kmer[i] = ~(kmer[i]);
  }
  kmer[lenQ] &= ((((KmerUnit)0x1) << (lenR << 1)) - 1);

  int32_t currHQ = lenQ;
  int32_t currHR = lenR - 1;
  int32_t currLQ = 0;
  int32_t currLR = 0;

  while (currHQ != currLQ || currHR != currLR)
  {
    uint64_t currHN = (kmer[currHQ] >> (currHR << 1)) & 0x3;
    uint64_t currLN = (kmer[currLQ] >> (currLR << 1)) & 0x3;
    kmer[currHQ] = (kmer[currHQ] & (~((KmerUnit)0x3 << (currHR << 1)))) + (currLN << (currHR << 1));
    kmer[currLQ] = (kmer[currLQ] & (~((KmerUnit)0x3 << (currLR << 1)))) + (currHN << (currLR << 1));

    if (currHR == 0)
    {
      currHR = nNtPerUnit - 1;
      currHQ--;
    }
    else
    {
      currHR--;
    }
    if (currLR == nNtPerUnit - 1)
    {
      currLR = 0;
      currLQ++;
    }
    else
    {
      currLR++;
    }
  }
}


bool Kmer::operator==(const Kmer & dst) const
{
  // check static variables
  for (int32_t i = 0; i <= lenQ; i++)
  {
    if (dst.kmer[i] != kmer[i]) return false;
  }
  return true;
}


bool Kmer::operator<(const Kmer & dst) const
{
  // check static variables
  for (int32_t i = lenQ; i >= 0; i--)
  {
    if (dst.kmer[i] > kmer[i]) return true;
    if (dst.kmer[i] < kmer[i]) return false;
  }
  return false;
}

#endif


nucleotide_type Kmer::charToNucleotide(char n)
{
  switch (n)
  {
    case 'A': return nt_A;
    case 'T': return nt_T;
    case 'G': return nt_G;
    case 'C': return nt_C;
    default:  return nt_A;
  }
}


nucleotide_type Kmer::charToComplementaryNucleotide(char n)
{
  switch (n)
  {
    case 'A': return nt_T;
    case 'T': return nt_A;
    case 'G': return nt_C;
    case 'C': return nt_G;
    default:  return nt_T;
  }
}


bool Kmer::isValidNucleotide(char n)
{
  return (n == 'A' || n == 'T' || n == 'G' || n == 'C');
}


char Kmer::nucleotideToChar(nucleotide_type n)
{
  switch (n)
  {
    case nt_A: return 'A';
    case nt_T: return 'T';
    case nt_G: return 'G';
    case nt_C: return 'C';
    default:   return 'I';
  }
}

